unit uSBDUnit_MemoryManagement;
interface


function  UseMemoryMetering: boolean;
procedure SuppressMemoryMetering( DeltaSuppressionCount: integer);
function  IsMeteringSuppressed: boolean;
function  ComputeNettAllocation: int64;

implementation



uses Windows;
var
  UseMemoryMeteringVar: boolean = True;
  Original_MM: TMemoryManagerEx;
  Metered_MM : TMemoryManagerEx;
  ExpectedMemoryLeakers: array[0..255] of pointer;
  ExpectedLeakerCursor: cardinal;
  NettAllocation: int64;
  MemoryGate: TRTLCriticalSection;
  SuppressionCount: integer;

function Metered_GetMem( Size: Integer): Pointer;                     forward;
function Metered_FreeMem( P: Pointer): Integer;                       forward;
function Metered_ReallocMem( P: Pointer; Size: Integer): Pointer;     forward;
function Metered_AllocMem( Size: Cardinal): Pointer;                  forward;
function Metered_RegisterExpectedMemoryLeak( P: Pointer): Boolean;    forward;
function Metered_UnregisterExpectedMemoryLeak( P: Pointer): Boolean;  forward;


function  UseMemoryMetering: boolean;
begin
result := UseMemoryMeteringVar
end;


procedure SuppressMemoryMetering( DeltaSuppressionCount: integer);
begin
Inc( SuppressionCount, DeltaSuppressionCount)
end;


function Leakage( const MS1, MS2: TMemoryManagerState): integer;
var
  j: integer;
  SMBSize1, SMBSize2: Int64;
begin
SMBSize1 := 0;
SMBSize2 := 0;
for j := 0 to NumSmallBlockTypes - 1 do
  begin
  Inc( SMBSize1, (MS1.SmallBlockTypeStates[j].InternalBlockSize * MS1.SmallBlockTypeStates[j].AllocatedBlockCount));
  Inc( SMBSize2, (MS2.SmallBlockTypeStates[j].InternalBlockSize * MS2.SmallBlockTypeStates[j].AllocatedBlockCount))
  end;
result := SMBSize2 - SMBSize1 +
    (Int64(MS2.TotalAllocatedMediumBlockSize) - Int64(MS1.TotalAllocatedMediumBlockSize)) +
    (Int64(MS2.TotalAllocatedLargeBlockSize ) - Int64(MS1.TotalAllocatedLargeBlockSize ))
end;


{$WARNINGS OFF}
procedure InitUnit_MemoryMangement;
var
  MS1, MS2, MS3: TMemoryManagerState;
  CanUseMemoryStateManager: boolean;
  Pntr: pointer;
  Size: integer;
begin
NettAllocation := 0;
ExpectedLeakerCursor := 0;
FillChar( MemoryGate, SizeOf( MemoryGate), 0);
FillChar( ExpectedMemoryLeakers, SizeOf( ExpectedMemoryLeakers), 0);
System.GetMemoryManager( Original_MM);
Metered_MM.GetMem     :=  Metered_GetMem;
Metered_MM.FreeMem    :=  Metered_FreeMem;
Metered_MM.ReallocMem :=  Metered_ReallocMem;
Metered_MM.AllocMem   :=  Metered_AllocMem;
Metered_MM.RegisterExpectedMemoryLeak   := Metered_RegisterExpectedMemoryLeak;
Metered_MM.UnregisterExpectedMemoryLeak := Metered_UnregisterExpectedMemoryLeak;

// Test. Is is safe to use the Embarcadero MemoryManagerState ?
Size := 1234;
GetMemoryManagerState( MS1);
GetMem( Pntr, Size);
GetMemoryManagerState( MS2);
FreeMem( Pntr, Size);
GetMemoryManagerState( MS3);
CanUseMemoryStateManager := (Leakage( MS1, MS3) = 0) and
                            (Leakage( MS1, MS2) >= Size);

// If not, use metering.
UseMemoryMeteringVar := not CanUseMemoryStateManager;
if not UseMemoryMeteringVar then exit;
InitializeCriticalSection( MemoryGate);
System.SetMemoryManager( Metered_MM);
end;


procedure DoneUnit_MemoryMangement;
begin
if not UseMemoryMeteringVar then exit;
System.SetMemoryManager( Original_MM);
DeleteCriticalSection( MemoryGate)
end;
{$WARNINGS ON}


procedure Meter( SizeOut: integer); inline;
begin
if SuppressionCount > 0 then exit;
{$IFNDEF WIN64}
EnterCriticalSection( MemoryGate);
{$ENDIF}
Inc( NettAllocation, SizeOut);
{$IFNDEF WIN64}
LeaveCriticalSection( MemoryGate);
{$ENDIF}
end;


function Metered_GetMem( Size: Integer): Pointer;
begin
if Size <= 0 then
    result := nil
  else
    begin
    result := Original_MM.GetMem( Size + SizeOf( Integer));
    PInteger( result)^ := Size;
    Inc( PByte( result), SizeOf( Integer));
    Meter( Size)
    end
end;

function Metered_FreeMem( P: Pointer): Integer;
var
  Size: integer;
begin
if not assigned( P) then
    result := 0
  else
    begin
    Dec( PByte( P), SizeOf( Integer));
    Size := PInteger( P)^;
    result := Original_MM.FreeMem( P);
    Meter( -Size)
    end;
end;


function Metered_ReallocMem( P: Pointer; Size: Integer): Pointer;
var
  OldSize: integer;
begin
if assigned( P) then
    begin
    Dec( PByte( P), SizeOf( Integer));
    OldSize := PInteger( P)^
    end
  else
    OldSize := 0;
if Size >= 0 then
    begin
    result  := Original_MM.ReallocMem( P, Size + SizeOf( Integer));
    PInteger( result)^ := Size;
    Inc( PByte( result), SizeOf( Integer))
    end
  else
    result  := Original_MM.ReallocMem( P, 0);
Meter( Size - OldSize)
end;

function Metered_AllocMem( Size: Cardinal): Pointer;
begin
if Size = 0 then
    result := nil
  else
    begin
    result := Original_MM.AllocMem( Size + SizeOf( Integer));
    PInteger( result)^ := Size;
    Inc( PByte( result), SizeOf( Integer));
    Meter( Size)
    end
end;

function Metered_RegisterExpectedMemoryLeak( P: Pointer): Boolean;
begin
result := assigned( P);
if not result then exit;
Original_MM.RegisterExpectedMemoryLeak( P);
if SuppressionCount > 0 then exit;
Dec( PByte( P), SizeOf( Integer));
EnterCriticalSection( MemoryGate);
try
  result := ExpectedLeakerCursor <= High( ExpectedMemoryLeakers);
  if not result then exit;
  ExpectedMemoryLeakers[ExpectedLeakerCursor] := P;
  Inc( ExpectedLeakerCursor)
finally
  LeaveCriticalSection( MemoryGate);
  end
end;

function Metered_UnregisterExpectedMemoryLeak( P: Pointer): Boolean;
var
  j, CountToShuffle: integer;
begin
result := assigned( P);
if not result then exit;
Original_MM.UnregisterExpectedMemoryLeak( P);
Dec( PByte( P), SizeOf( Integer));
result := False;
EnterCriticalSection( MemoryGate);
try
  for j := ExpectedLeakerCursor - 1 downto 0 do
    begin
    result := ExpectedMemoryLeakers[j] = P;
    if not result then continue;
{$WARNINGS OFF}
    CountToShuffle := ExpectedLeakerCursor - j - 1;
{$WARNINGS ON}
    if CountToShuffle > 0 then
      Move( ExpectedMemoryLeakers[j+1], ExpectedMemoryLeakers[j], CountToShuffle * SizeOf( pointer));
    Dec( ExpectedLeakerCursor);
    if ExpectedLeakerCursor >= 0 then
      ExpectedMemoryLeakers[ ExpectedLeakerCursor] := nil;
    break
    end
finally
  LeaveCriticalSection( MemoryGate)
  end;
if SuppressionCount > 0 then
  result := True
end;


function ComputeNettAllocation: int64;
var
  j: integer;
begin
result := 0;
if not UseMemoryMeteringVar then exit;
EnterCriticalSection( MemoryGate);
try
  result := NettAllocation;
  if ExpectedLeakerCursor >= 1 then
    for j := 0 to ExpectedLeakerCursor - 1 do
      Dec( result, PInteger( ExpectedMemoryLeakers[j])^)
finally
  LeaveCriticalSection( MemoryGate)
  end
end;


function IsMeteringSuppressed: boolean;
begin
result := (not UseMemoryMeteringVar) or (SuppressionCount > 0)
end;

initialization
InitUnit_MemoryMangement;

finalization
DoneUnit_MemoryMangement;

end.
